from tianshou.env import ShmemVectorEnv, DummyVectorEnv, SubprocVectorEnv

from ship_envs.environment.imazu_problem_env_factory import new_case_env_by_index
import numpy as np


def get_single_env_by_task(task_name):
    if task_name == 'random':
        cases = list(range(1, 23))
        cases = np.random.choice(cases, 1, replace=False)
        return new_case_env_by_index(cases[0])
    elif 'case' in task_name:
        cases = task_name.split('_')
        cases = cases[1:]
        cases = list(map(int, cases))
        cases = np.random.choice(cases, 1, replace=False)
        return new_case_env_by_index(cases[0])
    elif 'task1' == task_name:
        return new_case_env_by_index(np.random.randint(1, 3))
        # return new_case_env_by_index(1)


def make_case_env_wrapper(index):
    # 使用闭包保存变量index
    def inner_make_case_env():
        # print("inner_make_case_env:", index)
        return new_case_env_by_index(index)

    return inner_make_case_env


def get_case_envs(env_num, case_index=None, use_mul_pro=True):
    if case_index is None:
        case_index = [1]
    envs = []
    env_indexs = np.random.choice(case_index, env_num)
    for index in env_indexs:
        envs.append(make_case_env_wrapper(index))
    if use_mul_pro:
        return SubprocVectorEnv(envs)
    else:
        return DummyVectorEnv(envs)


def get_random_envs(env_num, use_mul_pro=True):
    indexs = list(range(1, 23))
    return get_case_envs(env_num, indexs, use_mul_pro)


def get_case_envs_by_task(task_name, nums, use_mul_pro=True):
    if task_name == 'random':
        return get_random_envs(nums, use_mul_pro)
    elif 'case' in task_name:
        cases = task_name.split('_')
        cases = cases[1:]
        cases = list(map(int, cases))
        return get_case_envs(nums, cases, use_mul_pro)
    elif task_name == 'task1':
        envs = []
        single_env_num = nums // 2
        for index in range(1, 3):
            for _ in range(single_env_num):
                envs.append(make_case_env_wrapper(index))
        if use_mul_pro:
            return SubprocVectorEnv(envs)
        else:
            return DummyVectorEnv(envs)
